package com.neuedu.framework.config;

import com.neuedu.common.utils.Threads;

import java.util.Optional;
import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.ScheduledThreadPoolExecutor;
import java.util.concurrent.ThreadPoolExecutor;

import lombok.extern.slf4j.Slf4j;

import org.apache.commons.lang3.concurrent.BasicThreadFactory;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.core.task.AsyncTaskExecutor;
import org.springframework.core.task.TaskDecorator;
import org.springframework.scheduling.TaskScheduler;
import org.springframework.scheduling.concurrent.ThreadPoolTaskExecutor;
import org.springframework.scheduling.concurrent.ThreadPoolTaskScheduler;
import org.springframework.security.core.context.SecurityContext;
import org.springframework.security.core.context.SecurityContextHolder;
import org.springframework.web.context.request.RequestAttributes;
import org.springframework.web.context.request.RequestContextHolder;

/**
 * @author Punkhoo
 */
@Configuration
@Slf4j
public class TaskConfig {

    /**
     * 获取当前机器的核数, 不一定准确 请根据实际场景 CPU密集 || IO 密集
     */
    public static final int CPU_NUM = Runtime.getRuntime().availableProcessors();

    /**
     * 线程池名前缀
     */
    private static final String ASYNC_THREAD_NAME_PREFIX = "NeuEdu-Async-Task-Thread";

    /**
     * 线程池名前缀
     */
    private static final String SCHEDULE_THREAD_NAME_PREFIX = "NeuEdu-Schedule-Task-Thread";

    @Value("${thread.pool.corePoolSize:}")
    private Optional<Integer> corePoolSize;

    @Value("${thread.pool.maxPoolSize:}")
    private Optional<Integer> maxPoolSize;

    @Value("${thread.pool.queueCapacity:}")
    private Optional<Integer> queueCapacity;

    @Value("${thread.pool.awaitTerminationSeconds:}")
    private Optional<Integer> awaitTerminationSeconds;

    /**
     * 执行周期性或定时任务
     */
    @Bean(name = "scheduledExecutorService")
    protected ScheduledExecutorService scheduledExecutorService() {
        return new ScheduledThreadPoolExecutor(corePoolSize.orElse(CPU_NUM),
            new BasicThreadFactory.Builder().namingPattern(SCHEDULE_THREAD_NAME_PREFIX + "-%d").daemon(false).build(),
            new ThreadPoolExecutor.CallerRunsPolicy()) {
            @Override
            protected void afterExecute(Runnable r, Throwable t) {
                super.afterExecute(r, t);
                Threads.printException(r, t);
            }
        };
    }

    @Bean("taskScheduler")
    public TaskScheduler taskScheduler() {
        ThreadPoolTaskScheduler executor = new ThreadPoolTaskScheduler();
        // 核心线程大小 默认区 CPU 数量
        executor.setPoolSize(corePoolSize.orElse(CPU_NUM));
        executor.setThreadNamePrefix(SCHEDULE_THREAD_NAME_PREFIX);

        return executor;
    }

    @Bean(name = "asyncTaskExecutor")
    public AsyncTaskExecutor asyncTaskExecutor() {
        ThreadPoolTaskExecutor taskExecutor = new ThreadPoolTaskExecutor();
        // 核心线程大小 默认区 CPU 数量
        taskExecutor.setCorePoolSize(corePoolSize.orElse(CPU_NUM));
        // 最大线程大小 默认区 CPU * 2 数量
        taskExecutor.setMaxPoolSize(maxPoolSize.orElse(CPU_NUM * 2));
        // 队列最大容量
        taskExecutor.setQueueCapacity(queueCapacity.orElse(100));
        taskExecutor.setRejectedExecutionHandler(new ThreadPoolExecutor.CallerRunsPolicy());
        taskExecutor.setWaitForTasksToCompleteOnShutdown(true);
        taskExecutor.setAwaitTerminationSeconds(awaitTerminationSeconds.orElse(60));
        taskExecutor.setThreadNamePrefix(ASYNC_THREAD_NAME_PREFIX);

        taskExecutor.setTaskDecorator(securityContextTaskDecorator());

        taskExecutor.initialize();

        return taskExecutor;
    }

    @Bean
    public TaskDecorator securityContextTaskDecorator() {
        return new SecurityContextTaskDecorator();
    }

}

class SecurityContextTaskDecorator implements TaskDecorator {

    @Override
    public Runnable decorate(Runnable task) {
        SecurityContext securityContext = SecurityContextHolder.getContext();
        RequestAttributes requestAttributes = RequestContextHolder.currentRequestAttributes();

        return () -> {
            RequestContextHolder.setRequestAttributes(requestAttributes);
            SecurityContextHolder.setContext(securityContext);
            try {
                task.run();
            } finally {
                RequestContextHolder.resetRequestAttributes();
                SecurityContextHolder.clearContext();
            }
        };
    }

}